/*
 * Description: install exclusion handler
 */
#include "installation_scene_accelerator.h"
#include <Windows.h>
#include <wintrust.h>
#include <SoftPub.h>
#include <regex>
#include "performance/utils/perf_utils.h"
#include "performance/perf_constants.h"
#include "../process_monitor_injector.h"
#include "text_util/string_proc.h"
#include "common/path_utils.h"
#include "log.h"
#include "process_load_statistic/sys_cpu_calc.h"
#include "performance/utils/defender_utils.h"
#pragma comment (lib, "wintrust")

namespace PERF {

const char* DEFAULTINSTALLWHITELIST = "C:\\ProgramData\\VMAgent\\config\\InstallationList.json";
const std::wstring ALTPATTERNPREFIX = L"(app_\\d+.*)?";
const DWORD WaitingTime = 300000;

void InstallPackageItem::Clear()
{
    name = "";
    fileNamePattern = "";
    exclusiveProcesses.clear();
    exclusiveFolders.clear();
}
InstallationSceneAcclerator::InstallationSceneAcclerator()
{
}
void InstallationSceneAcclerator::Init()
{
    HANDLE hUserToken = NULL;
    if (PERF_UTILS.GetUserToken(&hUserToken)) {
        notAllowedFolder_.insert(COMMON::PATHUTILS::GetCommonAppDataFolder(hUserToken));
        notAllowedFolder_.insert(COMMON::PATHUTILS::GetUserAppDataLocalFolder(hUserToken));
        notAllowedFolder_.insert(COMMON::PATHUTILS::GetUserAppDataRoamingFolder(hUserToken));
        notAllowedFolder_.insert(COMMON::PATHUTILS::GetSystemFolder());
        CloseHandle(hUserToken);
    }
}
void InstallationSceneAcclerator::Quit()
{
    // Remove all added exclusives
    LOG_INFO("[INSTALL - Release] - begin");
    std::unique_lock<std::mutex> lock(mutex_);
    for (auto  iter : runningProcesses_) {
        RunningProcess& process = iter.second;
        for (auto i : process.exclusiveFolders) {
            DefenderUtils::RemoveExclusionItem(i, PERF::ExclusionType::PATH);
        }
        for (auto i : process.exclusiveProcesses) {
            DefenderUtils::RemoveExclusionItem(i, PERF::ExclusionType::PROCESS);
        }
    }
    runningProcesses_.clear();
    LOG_INFO("[INSTALL - Release] - end");
}
bool InstallationSceneAcclerator::OnBegin(std::shared_ptr<PerfInfo> perfInfo)
{
    if (perfInfo == nullptr || perfInfo->processId == 0) {
        return false;
    }
    InstallPackageItem appConfig {};
    if (perfInfo->scene == PERF::PERF_SCENE::APP_INSTALL && IsInstallScene_(perfInfo, appConfig)) {
        // add process itself to Defender ExclusionProcess
        AddExclusiveItems_(perfInfo, appConfig);
        std::thread tWaitingCycle([this, perfInfo]() {
            WaitForProcessExit_(perfInfo);
            });
        tWaitingCycle.detach();
    } else if (perfInfo->scene == PERF::PERF_SCENE::APP_CREATE_FOLDER) {
        LOG_INFO("[INSTALL - OnBegin] - start {}, {}({}) {}.", \
            perfInfo->sceneStr.c_str(), \
            perfInfo->appName.c_str(), \
            perfInfo->processId, \
            perfInfo->workFolder.c_str());
        AddExclusiveWorkFolder_(perfInfo, perfInfo->workFolder);
    }
    return true;
}
bool InstallationSceneAcclerator::OnEnd(std::shared_ptr<PerfInfo> perfInfo)
{
    if (perfInfo != nullptr && perfInfo->scene == PERF::PERF_SCENE::APP_INSTALL) {
        Stop_(perfInfo);
    } else {
        LOG_WARN("[INSTALL - OnEnd] - end {} failed, scene {} is not expected here.", \
            perfInfo->appName.c_str(), \
            perfInfo->sceneStr.c_str());
    }
    return true;
}
bool InstallationSceneAcclerator::IsInstallScene_(std::shared_ptr<PerfInfo> perfInfo, InstallPackageItem& appConfig)
{
    LOG_DEBUG("[INSTALL - IsInstallScene_] begin :{} {}", perfInfo->appName, perfInfo->appPath);
    if (perfInfo->scene == PERF::PERF_SCENE::APP_INSTALL && \
        APPWHITELIST.SearchInWhiteList(perfInfo->appName, appConfig)) {
        std::wstring dosPath = COMMON::PATHUTILS::GetProcessFilePath(perfInfo->processId);
        if (APPWHITELIST.IsSigned(dosPath)) {
            LOG_INFO("[INSTALL - IsInstallScene_] is install scene :{}, signature is verified", \
                perfInfo->appName);
            return true;
        } else {
            LOG_INFO("[INSTALL - IsInstallScene_] is not install scene :{}, signature is not verified ", \
                perfInfo->appName);
            return false;
        }
    }
    return false;
}
bool InstallationSceneAcclerator::Stop_(std::shared_ptr<PerfInfo> perfInfo)
{
    if (perfInfo == nullptr || perfInfo->scene != PERF::PERF_SCENE::APP_INSTALL) {
        return false;
    }
    std::unique_lock<std::mutex> lock(mutex_);
    if (runningProcesses_.count(perfInfo->appName) == 0) {
        return false;
    }
    std::set<DWORD>& processIds = runningProcesses_[perfInfo->appName].processId;
    if (processIds.count(perfInfo->processId) == 0) {
        return false;
    }
    processIds.erase(perfInfo->processId);
    if (processIds.size() == 0) {
        RemoveExclusiveItems_(perfInfo);
        runningProcesses_.erase(perfInfo->appName);
    } else {
        LOG_INFO("[INSTALL - Stop_] {} has other process running", perfInfo->appName);
    }
    LOG_INFO("[INSTALL - Stop_] {}({}) Stopped. All exclusion has been cleared", \
        perfInfo->appName, \
        perfInfo->processId);
    return true;
}
void InstallationSceneAcclerator::WaitForProcessExit_(std::shared_ptr<PerfInfo> perfInfo)
{
    if (perfInfo == nullptr) {
        return;
    }
    LOG_INFO("[INSTALL - WaitForProcessExit_] - begin {}({})", \
        perfInfo->appName.c_str(), \
        perfInfo->processId);
    if (perfInfo->processId != 0) {
        HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | SYNCHRONIZE, FALSE, perfInfo->processId);
        if (hProcess != INVALID_HANDLE_VALUE) {
            DWORD waitCode = WaitForSingleObject(hProcess, WaitingTime);
            CloseHandle(hProcess);
            if (waitCode == WAIT_TIMEOUT) {
                LOG_WARN("[INSTALL - WaitForProcessExit_] - Install Scene {} timeout end", \
                    perfInfo->appName.c_str());
            }
        } else {
            LOG_ERR("[INSTALL - WaitForProcessExit_] - Install Scene end:unable to obtain installation process");
        }
        Stop_(perfInfo);
    }
    LOG_INFO("[INSTALL - WaitForProcessExit_] - end {}({})", \
        perfInfo->appName.c_str(), \
        perfInfo->processId);
}
bool InstallationSceneAcclerator::AddExclusiveWorkFolder_(std::shared_ptr<PerfInfo> perfInfo, std::string& folder)
{
    if (perfInfo == nullptr) {
        return false;
    }
    std::unique_lock<std::mutex> lock(mutex_);
    if (runningProcesses_.count(perfInfo->appName) == 0) {
        LOG_ERR("[INSTALL - add exclusive folder] failed {}", perfInfo->appName.c_str());
    }
    if (notAllowedFolder_.count(folder)) {
        return false;
    }
    RunningProcess& process = runningProcesses_[perfInfo->appName];
    if (PERF::DefenderUtils::AddExclusionItem(folder, PERF::ExclusionType::PATH)) {
        process.exclusiveFolders.insert(folder);
        return true;
    } else {
        LOG_ERR("Install Scene:{} add exclusive folder {} failed", perfInfo->appName.c_str(), folder);
    }
    return false;
}
bool InstallationSceneAcclerator::AddExclusiveItems_(std::shared_ptr<PerfInfo> perfInfo, const InstallPackageItem& item)
{
    if (perfInfo == nullptr) {
        return false;
    }
    std::unique_lock<std::mutex> lock(mutex_);
    LOG_DEBUG("[INSTALL - AddExclusiveItems_ begin] {}({})", perfInfo->appName, perfInfo->processId);
    if (runningProcesses_.count(perfInfo->appName) == 0) {
        LOG_DEBUG("[INSTALL - AddExclusiveItems_] new {}", perfInfo->appName);
    }
    RunningProcess& process = runningProcesses_[perfInfo->appName];
    if (process.processId.size() > 0) {
        process.processId.insert(perfInfo->processId);
        LOG_DEBUG("[INSTALL - AddExclusiveItems_ skip] {} has {} processes running", \
            perfInfo->appName, \
            process.processId.size());
        return true;
    }
    PERF::DefenderUtils::AddExclusionItem(perfInfo->appName, PERF::ExclusionType::PROCESS);
    process.exclusiveProcesses.insert(perfInfo->appName);
    process.processId.insert(perfInfo->processId);
    bool ret = true;
    for (auto i : item.exclusiveFolders) {
        if (PERF::DefenderUtils::AddExclusionItem(i, PERF::ExclusionType::PATH)) {
            process.exclusiveFolders.insert(i);
        } else {
            ret = false;
        }
    }
    for (auto i : item.exclusiveProcesses) {
        if (PERF::DefenderUtils::AddExclusionItem(i, PERF::ExclusionType::PROCESS)) {
            process.exclusiveProcesses.insert(i);
        } else {
            ret = false;
        }
    }
    LOG_INFO("Install Scene:{} AddExclusiveItems_ end {} {} {}", \
        perfInfo->appName, \
        process.exclusiveProcesses.size(), \
        process.exclusiveFolders.size());
    return ret;
}
bool InstallationSceneAcclerator::RemoveExclusiveItems_(std::shared_ptr<PerfInfo> perfInfo)
{
    if (perfInfo == nullptr) {
        return false;
    }
    if (runningProcesses_.count(perfInfo->appName) == 0) {
        return false;
    }
    RunningProcess& process = runningProcesses_[perfInfo->appName];
    bool ret = true;

    LOG_INFO("[INSTALL - RemoveExclusiveItems_ begin]:{} remove exclusive items {},{}", \
        perfInfo->appName, \
        (int)process.exclusiveFolders.size(), \
        (int)process.exclusiveProcesses.size());

    for (auto i : process.exclusiveFolders) {
        ret &= DefenderUtils::RemoveExclusionItem(i, PERF::ExclusionType::PATH);
    }
    for (auto i : process.exclusiveProcesses) {
        ret &= DefenderUtils::RemoveExclusionItem(i, PERF::ExclusionType::PROCESS);
    }
    LOG_INFO("[INSTALL - RemoveExclusiveItems_ end]:{} remove exclusive items", perfInfo->appName);
    return ret;
}

// Begin implementation of AppWhiteList

void AppWhiteList::Init()
{
    const char* CSIDLAPPDATAHOLDER = "%CSIDL_APPDATA%";
    const char* CSIDLLOCALAPPDATA_HOLDER = "%CSIDL_LOCAL_APPDATA%";
    const char* CSIDLCOMMONAPPDATA_HOLDER = "%CSIDL_COMMON_APPDATA%";
    const char* CSIDLCOMMONDOCUMENTSHOLDER = "CSIDL_LOCALLOW_APPDATA";
    const char* CSIDLPROGRAMFILESHOLDER = "%CSIDL_PROGRAM_FILES%";
    const char* CSIDLPROGRAMFILESHOLDER_86 = "%CSIDL_PROGRAM_FILESX86%";
    const char* CSIDLLOCALAPPDATALOW_HOLDER = "%CSIDL_LOCALLOW_APPDATA%";

    HANDLE hUserToken = NULL;
    if (!getUserToken_.load() && PERF_UTILS.GetUserToken(&hUserToken)) {
        std::lock_guard<std::mutex> lock(mutex_);
        userPathHolderMap_.insert(std::make_pair<std::string, std::string>(CSIDLAPPDATAHOLDER, \
            COMMON::PATHUTILS::GetUserAppDataRoamingFolder(hUserToken)));
        userPathHolderMap_.insert(std::make_pair<std::string, std::string>(CSIDLLOCALAPPDATA_HOLDER, \
            COMMON::PATHUTILS::GetUserAppDataLocalFolder(hUserToken)));
        userPathHolderMap_.insert(std::make_pair<std::string, std::string>(CSIDLCOMMONAPPDATA_HOLDER, \
            COMMON::PATHUTILS::GetCommonAppDataFolder(hUserToken)));
        userPathHolderMap_.insert(std::make_pair<std::string, std::string>(CSIDLCOMMONDOCUMENTSHOLDER, \
            COMMON::PATHUTILS::GetPublicDocumentsFolder(hUserToken)));
        userPathHolderMap_.insert(std::make_pair<std::string, std::string>(CSIDLPROGRAMFILESHOLDER, \
            COMMON::PATHUTILS::GetProgramFilesFolder(hUserToken)));
        userPathHolderMap_.insert(std::make_pair<std::string, std::string>(CSIDLPROGRAMFILESHOLDER_86, \
            COMMON::PATHUTILS::GetProgramFilesX86Folder(hUserToken)));
        userPathHolderMap_.insert(std::make_pair<std::string, std::string>(CSIDLLOCALAPPDATALOW_HOLDER, \
            COMMON::PATHUTILS::GetUserAppDataLocalLowFolder(hUserToken)));
        CloseHandle(hUserToken);
        ReadAppWhiteList_();
        getUserToken_.store(true);
        LOG_INFO("[INSTALL-START] - read config file ,add {} software", appWhiteList_.size());
    }
}
bool AppWhiteList::SearchInWhiteList(const std::string& appName, InstallPackageItem& item)
{
    Init();
    for (size_t i = 0; i < appNameRegex_.size(); i++) {
        if (std::regex_match(StringProc::StringToWString(appName), appNameRegex_[i])) {
            item = appWhiteList_[i];
            return true;
        }
    }
    return false;
}
bool AppWhiteList::ReadConfigFile_(std::string& config)
{
    HANDLE hConfigFile = CreateFile(DEFAULTINSTALLWHITELIST, GENERIC_READ, 0, nullptr, OPEN_EXISTING,
        FILE_ATTRIBUTE_NORMAL, nullptr);
    if (hConfigFile == INVALID_HANDLE_VALUE) {
        LOG_ERR("[INSTALL - ReadConfigFile_] - error occurred, file handle is nullptr, error code: {}",
            GetLastError());
        return false;
    }
    DWORD size = GetFileSize(hConfigFile, nullptr);
    DWORD sizeReaded = 0;
    char* data = new (std::nothrow) char[size + 1];
    if (data == nullptr) {
        LOG_ERR("[INSTALL - ReadConfigFile_] - error occurred, allocate memory failed!");
        if (hConfigFile != INVALID_HANDLE_VALUE) {
            CloseHandle(hConfigFile);
            hConfigFile = INVALID_HANDLE_VALUE;
        }
        return false;
    }
    SecureZeroMemory(data, size + 1);
    if (!ReadFile(hConfigFile, data, size, &sizeReaded, nullptr) || sizeReaded != size) {
        LOG_ERR("[INSTALL - ReadConfigFile_] - error occurred, reading file failed, error code: {}", GetLastError());
        if (hConfigFile != INVALID_HANDLE_VALUE) {
            CloseHandle(hConfigFile);
            hConfigFile = INVALID_HANDLE_VALUE;
        }
        delete[] data;
        return false;
    }
    if (hConfigFile != INVALID_HANDLE_VALUE) {
        CloseHandle(hConfigFile);
        hConfigFile = INVALID_HANDLE_VALUE;
    }
    config = std::string(data);
    delete[] data;
    return true;
}
bool AppWhiteList::ReadAppWhiteList_()
{
    appWhiteList_.clear();
    std::string config;
    if (!ReadConfigFile_(config)) {
        return false;
    }
    ParseAppWhiteList_(config.c_str());
    return true;
}
bool AppWhiteList::ParseAppWhiteList_(const char* data)
{
    cJSON* root = cJSON_Parse(data);
    if (!root) {
        return false;
    }
    cJSON* config = cJSON_GetObjectItem(root, "config");
    if (!config) {
        cJSON_Delete(root);
        return false;
    }
    cJSON* softwareList = cJSON_GetObjectItem(config, "SoftWareList");
    if (!softwareList) {
        cJSON_Delete(root);
        return false;
    }
    int listSize = cJSON_GetArraySize(softwareList);
    InstallPackageItem software;
    for (int i = 0; i < listSize; i++) {
        cJSON* item = cJSON_GetArrayItem(softwareList, i);
        if (item && ParseApp_(item, software)) {
            try {
                std::wstring wpattern = ALTPATTERNPREFIX + StringProc::StringToWStringByUTF8(software.fileNamePattern);
                std::wregex reg_pattern(wpattern, std::regex_constants::icase);
                appNameRegex_.emplace_back(reg_pattern);
                appWhiteList_.emplace_back(software);
            }
            catch (...) {
                LOG_ERR("Install::Construct regex error : {} {}", software.name, software.fileNamePattern);
            }
        }
    }
    cJSON_Delete(root);
    return true;
}
void AppWhiteList::ParseExclusionItems(cJSON* folder, std::vector<std::string>& vItems)
{
    int size = cJSON_GetArraySize(folder);
    for (int i = 0; i < size; i++) {
        cJSON* v = cJSON_GetArrayItem(folder, i);
        if (v && v->valuestring) {
            vItems.emplace_back(std::string(v->valuestring));
        }
    }
}
bool AppWhiteList::ParseApp_(cJSON* item, InstallPackageItem& parsed)
{
    parsed.Clear();
    cJSON* name = cJSON_GetObjectItem(item, "SName");
    if (name && name->valuestring) {
        parsed.name = std::string(name->valuestring);
    } else {
        return false;
    }
    cJSON* fileNamePattern = cJSON_GetObjectItem(item, "Pattern");
    if (fileNamePattern && fileNamePattern->valuestring) {
        parsed.fileNamePattern = std::string(fileNamePattern->valuestring);
    } else {
        return false;
    }
    cJSON* process = cJSON_GetObjectItem(item, "Process");
    if (process != nullptr)
        ParseExclusionItems(process, parsed.exclusiveProcesses);
    cJSON* folder = cJSON_GetObjectItem(item, "Folder");
    if (folder != nullptr) {
        ParseExclusionItems(folder, parsed.exclusiveFolders);
        for (auto& item : parsed.exclusiveFolders) {
            ReplaceUserPathHolder_(item);
        }
    }
    return true;
}
bool AppWhiteList::IsSigned(const std::wstring& filePath)
{
    WINTRUST_FILE_INFO fileData;
    SecureZeroMemory(&fileData, sizeof(fileData));
    
    fileData.cbStruct = sizeof(WINTRUST_FILE_INFO);
    fileData.pcwszFilePath = filePath.c_str();
    fileData.hFile = NULL;
    fileData.pgKnownSubject = NULL;

    GUID WVTPolicyGUID = WINTRUST_ACTION_GENERIC_VERIFY_V2;
    WINTRUST_DATA trustData;
    SecureZeroMemory(&trustData, sizeof(trustData));

    trustData.cbStruct = sizeof(trustData);
    trustData.pPolicyCallbackData = NULL;
    trustData.pSIPClientData = NULL;
    trustData.dwUIChoice = WTD_UI_NONE;
    trustData.fdwRevocationChecks = WTD_REVOKE_NONE;
    trustData.dwUnionChoice = WTD_CHOICE_FILE;
    trustData.dwStateAction = WTD_STATEACTION_VERIFY;
    trustData.hWVTStateData = NULL;
    trustData.pwszURLReference = NULL;
    trustData.dwUIContext = 0;
    trustData.pFile = &fileData;
    LONG lStatus = WinVerifyTrust(
        NULL,
        &WVTPolicyGUID,
        &trustData);

    bool isVerified = (lStatus == ERROR_SUCCESS);
    trustData.dwStateAction = WTD_STATEACTION_CLOSE;
    lStatus = WinVerifyTrust(
        NULL,
        &WVTPolicyGUID,
        &trustData);
    if (lStatus != ERROR_SUCCESS) {
        LOG_ERR("failed to close verify");
    }
    return isVerified;
}
void AppWhiteList::ReplaceUserPathHolder_(std::string& path)
{
    LOG_DEBUG("replace holder before{}", path);
    for (const auto& item : userPathHolderMap_) {
        if (path.find(item.first) == 0) {
            path.replace(0, item.first.length(), item.second);
            break;
        }
    }
    LOG_DEBUG("replace holder after{}", path);
}
// End implementation of AppWhiteList
}
// End of namespace PERF
